package com.jl;

import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

@Slf4j
public class ThreadPoolUtils {

    /**
     * 创建线程池
     *
     * @return
     */
    public static NewPoll newPoll() {
        return new NewPoll();
    }

    /**
     * 执行线程池
     *
     * @param threadName 线程名称
     * @param num        任务数
     * @return
     */
    public static Execute excPoll(String threadName, int num) {
        return new Execute(threadName, num);
    }

    public static class NewPoll {
        private ThreadPoolTaskExecutor threadPoolTaskExecutor = new ThreadPoolTaskExecutor();
        private List<String> list = new ArrayList<>();

        /**
         * 核心线程数
         *
         * @param coreSize
         * @return
         */
        public NewPoll coreSize(int coreSize) {
            threadPoolTaskExecutor.setCorePoolSize(coreSize);
            list.add("coreSize");
            return this;
        }

        /**
         * 最大线程数
         *
         * @param maxSize
         * @return
         */
        public NewPoll maxSize(int maxSize) {
            threadPoolTaskExecutor.setMaxPoolSize(maxSize);
            list.add("maxSize");
            return this;
        }

        /**
         * 队列大小
         *
         * @param queueSize
         * @return
         */
        public NewPoll queueSize(int queueSize) {
            threadPoolTaskExecutor.setQueueCapacity(queueSize);
            list.add("queueSize");
            return this;
        }

        /**
         * 线程空闲存活
         *
         * @param keepAlive
         * @return
         */
        public NewPoll keepAlive(int keepAlive) {
            threadPoolTaskExecutor.setKeepAliveSeconds(keepAlive);
            list.add("keepAlive");
            return this;
        }

        /**
         * 线程名称
         *
         * @param name
         * @return
         */
        public NewPoll name(String name) {
            threadPoolTaskExecutor.setThreadNamePrefix(name);
            list.add("name");
            return this;
        }

        /**
         * 拒绝策略
         *
         * @param rejectedExecutionHandler
         * @return
         */
        public NewPoll rejected(RejectedExecutionHandler rejectedExecutionHandler) {
            threadPoolTaskExecutor.setRejectedExecutionHandler(rejectedExecutionHandler);
            list.add("rejected");
            return this;
        }

        /**
         * 构造
         *
         * @return
         */
        public ThreadPoolTaskExecutor builder() {
            if (!list.contains("rejected")) {
                threadPoolTaskExecutor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
            }
            return threadPoolTaskExecutor;
        }
    }

    public static class Execute {

        private ThreadPoolTaskExecutor executor;

        private CountDownLatch latch;

        private boolean isError = false;

        public Execute(String threadName, int size) {
            this.executor = SpringBeanUtils.getBean(threadName, ThreadPoolTaskExecutor.class);
            this.latch = new CountDownLatch(size);
        }

        private Error error = e -> {
            throw new ExceptionInfo(e);
        };

        public static class ExceptionInfo extends RuntimeException {

            public ExceptionInfo(Throwable e) {
                super(e);
            }

        }

        /**
         * 线程执行
         *
         * @param function
         * @return
         */
        public Execute execute(ExecuteFunction function) {
            executor.execute(() -> {
                try {
                    function.exec();
                } catch (Exception e) {
                    if (isError) {
                        error.run(e);
                    } else {
                        log.error("", e);
                    }
                } finally {
                    latch.countDown();
                }
            });
            return this;
        }

        /**
         * 异常处理
         *
         * @param error
         * @return
         */
        public Execute error(Error error) {
            this.error = error;
            isError = true;
            return this;
        }

        /**
         * 任务创建完毕，阻塞等待
         */
        public void end() {
            try {
                latch.await();
            } catch (Exception e) {
                log.error("", e);
            }
        }

        /**
         * 任务创建完毕，阻塞等待
         *
         * @param seconds 最大等待时间，单位：s
         */
        public void end(int seconds) {
            try {
                latch.await(seconds, TimeUnit.SECONDS);
            } catch (Exception e) {
                log.error("", e);
            }
        }

        @FunctionalInterface
        public interface ExecuteFunction {
            void exec();
        }

        @FunctionalInterface
        public interface Error {
            void run(Throwable e);
        }
    }
}
